import mne
import matplotlib.pyplot as plt
from sklearn import preprocessing
import numpy as np
from keras.utils.np_utils import to_categorical
from sklearn.model_selection import train_test_split
from scipy import signal
"""
    Physionet MI-EEG Dataset
    64 channels EEG?160hz freq, 4 seconds MI-task
    14 runs for each of the 109 subjects
        runs [1, 2] is baseline
        others with marker
            T0   : rest, 
            T1/T2: left/right fist in runs [3, 4, 7, 8, 11, 12]
                   both fists/feet in runs [5, 6, 9, 10, 13, 14]

"""
import os
from scipy.fftpack import fft
import pandas as pd
data_path = "/media/brainseek/dataset/109subjects/LB_2s_250Hz/"
save_path = "/media/brainseek/dataset/109subjects/LB_2s_250Hz_80train_20test/"


LR_fist_run = [3, 4, 7, 8, 11, 12]


if __name__ == '__main__':
    subjects_bad = [38, 88, 89, 92, 100, 104]
    subjects_good = [n for n in np.arange(1, 110) if n not in subjects_bad]
    train_data_total  = np.empty((0, 320, 2))
    test_data_total   = np.empty((0, 320, 2))
    train_label_total = np.empty((0, 2))
    test_label_total  = np.empty((0, 2))
    if not os.path.exists(save_path):
        os.mkdir(save_path)
    for subs in subjects_good:
        if subs <= 80:
            train_data_raw = np.load(os.path.join(data_path, "train_data" + str(subs) +".npy"))
            test_data_raw  = np.load(os.path.join(data_path, "test_data"  + str(subs) +".npy"))
            train_label_raw = np.load(os.path.join(data_path,"train_label"+ str(subs) +".npy"))
            test_label_raw  = np.load(os.path.join(data_path,"test_label" + str(subs) +".npy"))

            train_data_total = np.concatenate((train_data_total, train_data_raw), axis=0)
            train_data_total = np.concatenate((train_data_total, test_data_raw), axis=0)
            train_label_total = np.concatenate((train_label_total, train_label_raw), axis=0)
            train_label_total = np.concatenate((train_label_total, test_label_raw), axis=0)

        else:
            train_data_raw = np.load(os.path.join(data_path, "train_data" + str(subs) +".npy"))
            test_data_raw  = np.load(os.path.join(data_path, "test_data"  + str(subs) +".npy"))
            train_label_raw = np.load(os.path.join(data_path,"train_label"+ str(subs) +".npy"))
            test_label_raw  = np.load(os.path.join(data_path,"test_label" + str(subs) +".npy"))

            test_data_total = np.concatenate((test_data_total, train_data_raw), axis=0)
            test_data_total = np.concatenate((test_data_total, test_data_raw), axis=0)
            test_label_total = np.concatenate((test_label_total, train_label_raw), axis=0)
            test_label_total = np.concatenate((test_label_total, test_label_raw), axis=0)
    train_data_total = signal.resample_poly(train_data_total, 500, 320, axis=1)
    test_data_total = signal.resample_poly(test_data_total, 500, 320, axis=1)
    # #huatu ceshi
    # for i in range(len(data)):
    #     plt.figure()
    #     plt.subplot(2, 1, 1)
    #     x = np.arange(data.shape[1])
    #     plt.plot(x, data[i, :, 0])
    #     x1 = np.arange(train_data_total.shape[1])
    #     fft_y = fft(train_data_total[i,:,0])
    #     abs_y = np.abs(fft_y)
    #     xf = np.arange(0,160,1)
    #     plt.subplot(2, 1, 2)
    #     plt.plot(x1,train_data_total[i,:,0])
    #     plt.show()
    # # huatu ceshi

    np.save(os.path.join(save_path, "train_data_total" ),  train_data_total, allow_pickle=True)
    np.save(os.path.join(save_path, "test_data_total" ),   test_data_total, allow_pickle=True)
    np.save(os.path.join(save_path, "train_label_total" ), train_label_total, allow_pickle=True)
    np.save(os.path.join(save_path, "test_label_total" ),  test_label_total, allow_pickle=True)
    print(train_data_total.shape)
    print(test_data_total.shape)
    print(train_label_total.shape)
    print(test_label_total.shape)